# -*- coding: utf-8 -*-
"""
mathmodeling_SAK.programming.analysis
=====================================

定义了对 Pyomo 规划模型输出结果进行分析的一些工具
"""

import numpy as np
from pyomo.environ import *

def pyomo_var2ndarray(model, var_name):
    """ 
    从 Pyomo 模型中提取变量，并按其原始多维索引结构转换为 numpy 数组。

    输入参数
    ----------
    model : pyomo.ConcreteModel
        通过 Pyomo 构建的规划模型对象
    var_name : str
        要从模型中提取的变量的名字，假定该变量已在模型中定义

    返回值
    ----------
    numpy.array
        多维数组，包含按多维索引排序的指定变量的值
        
    示例
    ----------
    
    >>> model = your_pyomo_model
    >>> ndarray = extract_pyomo_var_to_ndarray(model, 'P')
    >>> print(ndarray)
    """

    # 获取变量对象
    var = getattr(model, var_name)
    
    # 确保是索引变量
    if not var.is_indexed():
        return np.array([var.value])
    
    # 获取所有索引并计算数组维度
    all_indices = list(var.index_set())
    # 转换索引为0-based
    if all(isinstance(idx, tuple) for idx in all_indices):
        # 如果是多维索引，提取维度信息
        dimensions = np.max(
            np.array(all_indices) - 1, axis=0) + 1
    else:
        # 如果是单维索引，直接转换成一维数组
        dimensions = max(all_indices)

    # 创建一个多维数组并填充数据
    values_array = np.zeros(dimensions, dtype=np.float64)
    for idx in all_indices:
        if isinstance(idx, tuple):
            corrected_idx = tuple(x-1 for x in idx)  
            # 从1-based转为0-based索引
            values_array[corrected_idx] = var[idx].value
        else:
            values_array[idx - 1] = var[idx].value  
            # 从1-based转为0-based索引

    return values_array